'''
# 第三种融合方式 分别进行encoder, 最后全连接层融合三种特征
'''

import pandas as pd
import numpy as np
import os
import time
import datetime
import shutil
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, TerminateOnNaN
from tensorflow.keras.optimizers import SGD, Adam
from tensorflow.keras.regularizers import l2
from tensorflow.keras.initializers import glorot_uniform
from tensorflow.keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization,Conv2D
from tensorflow.keras.layers import AveragePooling2D, MaxPooling2D, Flatten, Concatenate, Dropout

model_name = "two_resnet50_ct_suv_info_3_"
time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
log_dir = "/home/zmy/keras_test/logs/" + model_name + time_str
work_dir = "/home/zmy/keras_test/workdir/" + model_name + time_str
model_save_path = "/home/zmy/keras_test/models/" + model_name
learn_rate = 0.0001
Batch_size = 4
Epoch = 100

def standard_scaler(features):
    scaler = MinMaxScaler()
    # scaler = StandardScaler()
    x_train = scaler.fit_transform(features)
    
    return x_train

def read_data():
    test_data = pd.read_csv('/data1/zmy/data/divide_csv/two/test.csv')
    train_data = pd.read_csv('/data1/zmy/data/divide_csv/two/train.csv')
    
    test_features = []
    test_labels = []
    test_ctslice = []
    test_suvslice = []
    
    train_features = []
    train_labels = []
    train_ctslice = []
    train_suvslice = []

    for i in range(len(test_data)):
        one_feature = [test_data['z'][i], test_data['x'][i]/512, test_data['y'][i]/512, test_data['r'][i],
                       test_data['patientWeight'][i], test_data['patientSex'][i], test_data['patientAge'][i],
                       test_data['patientSize'][i], test_data['local_suvmax'][i], test_data['local_suvmin'][i],
                       test_data['local_suvavg'][i], test_data['local_suvstd'][i], test_data['local_suvvar'][i]]

        test_features.append(one_feature)
        test_labels.append(int(test_data['cancer_type'][i]) - 1)
        ct_name = '/data1/zmy/data/Slice/two/test/CTSlice/'+str(test_data['patientID'][i])+'_'+str(test_data['cancer_type'][i])+'_CTSlice.npy'
        ct = np.load(ct_name)
        ct = ct[:, :, np.newaxis]
        test_ctslice.append(ct)
        suv_name = '/data1/zmy/data/Slice/two/test/SUVSlice/'+str(test_data['patientID'][i])+'_'+str(test_data['cancer_type'][i])+'_SUVSlice.npy'
        suv = np.load(suv_name)
        suv = suv[:, :, np.newaxis]
        test_suvslice.append(suv)

    for j in range(len(train_data)):
        one_feature = [train_data['z'][j], train_data['x'][j]/512, train_data['y'][j]/512, train_data['r'][j],
                       train_data['patientWeight'][j], train_data['patientSex'][j], train_data['patientAge'][j],
                       train_data['patientSize'][j], train_data['local_suvmax'][j], train_data['local_suvmin'][j],
                       train_data['local_suvavg'][j],  train_data['local_suvstd'][j], train_data['local_suvvar'][j]]
        train_features.append(one_feature)
        train_labels.append(int(train_data['cancer_type'][j]) - 1)
        ct_name = '/data1/zmy/data/Slice/two/train/CTSlice/'+str(train_data['patientID'][j])+'_'+str(train_data['cancer_type'][j])+'_CTSlice.npy'
        ct = np.load(ct_name)
        ct = ct[:, :, np.newaxis]
        train_ctslice.append(ct)
        suv_name = '/data1/zmy/data/Slice/two/train/SUVSlice/'+str(train_data['patientID'][j])+'_'+str(train_data['cancer_type'][j])+'_SUVSlice.npy'
        suv = np.load(suv_name)
        suv = suv[:, :, np.newaxis]
        train_suvslice.append(suv)

    x_train = np.asarray(train_features, dtype=np.float)
    x_train_ct = np.asarray(train_ctslice)
    x_train_suv = np.asarray(train_suvslice)
    y_train = np.asarray(train_labels, dtype=np.int)

    x_test = np.asarray(test_features, dtype=np.float)
    x_test_ct = np.asarray(test_ctslice)
    x_test_suv = np.asarray(test_suvslice)
    y_test = np.asarray(test_labels, dtype=np.int)

    x_train = standard_scaler(x_train)
    x_test = standard_scaler(x_test)


    # shuffle trainset
    m = x_train_ct.shape[0]
    permutation = list(np.random.permutation(m))
    x_train_ct = x_train_ct[permutation, :, :, :]
    x_train_suv = x_train_suv[permutation, :, :, :]
    x_train = x_train[permutation, :]
    y_train = y_train[permutation]

    return x_train, x_train_ct, x_train_suv, y_train, x_test, x_test_ct, x_test_suv, y_test


def identity_block(X, f, filters, stage, block):
    """
    Implementation of the identity block as defined in Figure 3

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network

    Returns:
    X -- output of the identity block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value. You'll need this later to add back to the main path.
    X_shortcut = X

    # First component of main path
    X = Conv2D(filters=F1, kernel_size=(1, 1), strides=(1, 1), padding='same', name=conv_name_base + '2a',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='same', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def convolutional_block(X, f, filters, stage, block, s=2):
    """
    Implementation of the convolutional block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network
    s -- Integer, specifying the stride to be used

    Returns:
    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value
    X_shortcut = X

    ##### MAIN PATH #####
    # First component of main path
    X = Conv2D(F1, (1, 1), strides=(s, s), padding='same', name=conv_name_base + '2a', kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='same', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    ##### SHORTCUT PATH #### (≈2 lines)
    X_shortcut = Conv2D(filters=F3, kernel_size=(1, 1), strides=(s, s), padding='same', name=conv_name_base + '1',
                        kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    X_shortcut = BatchNormalization(axis=3, name=bn_name_base + '1')(X_shortcut)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X

def base_ResNet50(X_input, index='one'):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    X
    """

    # Define the input as a tensor with shape input_shape
    # Stage 1
    X = Conv2D(filters=64, kernel_size=(7, 7), strides=(2, 2), padding='same', name='conv1'+index, kernel_initializer=glorot_uniform(seed=0))(X_input)
    X = BatchNormalization(axis=3, name='bn_conv1'+index)(X)
    X = Activation('relu')(X)
    X = MaxPooling2D((3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, f=3, filters=[64, 64, 256], stage=2, block='a'+index, s=1)
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='b'+index)
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='c'+index)

    ### START CODE HERE ###

    # Stage 3 (≈4 lines)
    X = convolutional_block(X, f=3, filters=[128, 128, 512], stage=3, block='a'+index, s=2)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='b'+index)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='c'+index)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='d'+index)

    # Stage 4 (≈6 lines)
    X = convolutional_block(X, f=3, filters=[256, 256, 1024], stage=4, block='a'+index, s=2)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='b'+index)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='c'+index)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='d'+index)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='e'+index)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='f'+index)

    # Stage 5 (≈3 lines)
    X = convolutional_block(X, f=3, filters=[512, 512, 2048], stage=5, block='a'+index, s=2)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='b'+index)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='c'+index)

    # AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    X = AveragePooling2D((2, 2), name="avg_pool"+index)(X)

    ### END CODE HERE ###

    # output layer
    X = Flatten()(X)
    return X

def ResNet50(main_input_shape=(512, 512, 1), aux_input_shape=(512, 512, 1), other_info_shape=(13, ), classes=1):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """

    # Define the input as a tensor with shape input_shape
    X_input = Input(main_input_shape)
    aux_X_input = Input(aux_input_shape)
    other_input = Input(other_info_shape)

    X = base_ResNet50(X_input, index='one')
    aux_X = base_ResNet50(aux_X_input, index='two')

    X = Concatenate()([X, aux_X, other_input])

    # X = Dense(1024, kernel_regularizer=l2(0.001), activation='relu')(X)
    # X = Dropout(0.5)(X)
    X = Dense(128, kernel_regularizer=l2(0.001), activation='relu', kernel_initializer=glorot_uniform(seed=0))(X)
    X = Dense(classes, activation='sigmoid', name='fc' + str(classes), kernel_initializer=glorot_uniform(seed=0))(X)

    # Create model
    model = Model(inputs=[X_input, aux_X_input, other_input], outputs=X, name=model_name)
    return model


def train():
    x_train, x_train_ct, x_train_suv, y_train, x_test, x_test_ct, x_test_suv, y_test = read_data()


    print("number of training examples = " + str(x_train.shape[0]))
    print("number of test examples = " + str(x_test.shape[0]))
    print("X_train shape: " + str(x_train.shape))
    print("X_train_ct shape: " + str(x_train_ct.shape))
    print("X_train_suv shape: " + str(x_train_suv.shape))
    print("Y_train shape: " + str(y_train.shape))

    print("X_test shape: " + str(x_test.shape))
    print("X_test_ct shape: " + str(x_test_ct.shape))
    print("X_test_suv shape: " + str(x_test_suv.shape))
    print("Y_test shape: " + str(y_test.shape))

    model = ResNet50()
    # adam = Adam(lr=learn_rate)
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    model.summary()

    if not os.path.exists(work_dir):
        os.makedirs(work_dir)
    # checkpoint = ModelCheckpoint(filepath=work_dir + "/" + model_name + "_" + "_e" + "{epoch:02d}-{val_loss:.4f}.hd5",
    #                              monitor='val_loss',
    #                              verbose=1,
    #                              save_best_only=False,
    #                              save_weights_only=False,
    #                              mode='auto',
    #                              period=1)
    # 每隔一轮且每当val_loss降低时保存一次模型
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    checkpoint_fixed_name = ModelCheckpoint(filepath=best_model_path,
                                            monitor='val_loss',
                                            verbose=1,
                                            save_best_only=True,
                                            save_weights_only=False,
                                            mode='auto',
                                            period=1)
    model.fit(x=[x_train_ct, x_train_suv, x_train],
              y=y_train,
              batch_size=Batch_size,
              epochs=Epoch,
              validation_data=([x_test_ct, x_test_suv, x_test], y_test),
              callbacks=[checkpoint_fixed_name, TensorBoard(log_dir=log_dir, write_graph=True, write_grads=True, write_images=True), TerminateOnNaN()]
              )


def predict():
    """
    进行预测
    """
    x_train, x_train_ct, x_train_suv, Y_train, x_test, x_test_ct, x_test_suv, Y_test = read_data()

    # 模型
    model_path = "/home/zmy/keras_test/models/two_resnet50_ct_suv_3/two_resnet50_ct_suv_3_best.hd5"
    model = ResNet50()
    model.load_weights(model_path, by_name=False)

    # 开始计时
    start_time = datetime.datetime.now()
    preds = model.predict([x_test_ct, x_test_suv, x_test])

    # 测试花费时间
    current_time = datetime.datetime.now()
    res = current_time - start_time
    print("Done in : ", res.total_seconds(), " seconds")
    print("model: " + model_path)
    print("label:", Y_test)
    print("predict:", preds)


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    if not os.path.exists(model_save_path):
        os.makedirs(model_save_path)
    train()
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    shutil.copy(best_model_path, model_save_path)

